#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jun 25 20:18:28 2021

@author: liangy
E-mail: liangy19@mails.tsinghua.edu.cn

For spherical lensing system

Unit System
The unit system we apply here is:
     distance         --> kiloparsec 
     angular distance --> arcsec
     mass             --> solar mass
     time             --> Gyr
     velocity         --> km/s
"""

import numpy as np
from scipy import constants
from scipy.optimize import root_scalar
from astropy.cosmology import FlatLambdaCDM
from sph_profile import sph_profile as sph

"""
The physical constants required
"""
G   = 4.301e-6          # gravitational constant in (km/s)^2*kpc/M_sun
h   = 0.6774            # reduced Hubble constant
H0  = 100*h             # Hubble constant in km/s/Mpc
Om0 = 0.3089            # matter density in cosmic inventory at z = 0
c   = constants.c*1e-3  # speed of light in km/s
arcsec2rad = np.pi/180/3600    # arcsec/rad

cosmo = FlatLambdaCDM(H0,Om0)

class sph_lensing():
    """
    To calculate the model angular Einstein radius in Bayesian framework
    """
    def __init__(self,lens_specification):
        self.lens_specification = lens_specification
        
        # if the cosmology is not determined, use the Planck2016
        try:
            self.cosmo = lens_specification["cosmo"]
        except:
            self.cosmo = cosmo
        
        # the redshift of deflector and source 
        self.zd  = lens_specification["zd"]    # the redshift of deflector
        self.zs  = lens_specification["zs"]    # the redshift of sources (float or np.ndarray)
        
        # angular diameter distances of sources, deflector in kpc
        self.Ds  = self.cosmo.angular_diameter_distance(self.zs).value*1e3
        self.Dd  = self.cosmo.angular_diameter_distance(self.zd).value*1e3
        self.Dds = self.cosmo.angular_diameter_distance_z1z2(self.zd,self.zs).value*1e3
        self.D_ratio = self.Dds/(self.Ds*self.Dd)
        self.sigma_crit = c**2*self.Ds/self.Dd/self.Dds/(4*np.pi*G)
        
        self.mass_model_list = lens_specification["mass_model_list"]
        self.mass_model_para = lens_specification["mass_model_para"]
        
        #self.sph_lens = sph(self.R_ein,dim = '2D')
        
    def get_einstein_radius(self,theta_ein):
        self.theta_ein = theta_ein
        return np.sqrt(4.0*G*self.D_ratio/c**2*self.projected_mass(theta_ein))/arcsec2rad
        
    def mock_einstein_radius(self,R_min = 0.05,R_max = 30.0):
        if self.lensing_equation(R_min/self.Dd/arcsec2rad)*self.lensing_equation(R_max/self.Dd/arcsec2rad)>0.0:
            self.theta_ein = -1.0
        else:
            lensing_sol = root_scalar(self.lensing_equation,
                                      bracket = (R_min/self.Dd/arcsec2rad,R_max/self.Dd/arcsec2rad))
            if lensing_sol.converged:
                self.theta_ein = lensing_sol.root
            else:
                self.theta_ein = -1.0
        return self.theta_ein
                
    def mock_secondary_parameter(self):
        if self.theta_ein > 0.0:
            Xi = 2*self.theta_ein*self.deflection_angle_d2(self.theta_ein)/(1-self.deflection_angle_d1(self.theta_ein))
        else:
            Xi = -100.0
        self.Xi = Xi
        return Xi
    
    def projected_mass(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in unit of radian
        xi = theta_rad*self.Dd          # radius in kpc
        proj_mass = 0.0
        sph_lens = sph(xi,dim="2D")
        for i,lens_mass in enumerate(self.mass_model_list):
            lens_mass_para_i = self.mass_model_para[i]
            if lens_mass == "power_law":
                proj_mass += sph_lens.M_power_law(*lens_mass_para_i)
            if lens_mass == "two_power_law":
                proj_mass += sph_lens.M_two_power_law(*lens_mass_para_i)
            if lens_mass == "sersic":
                proj_mass += sph_lens.M_sersic(*lens_mass_para_i)
            if lens_mass == "einasto":
                proj_mass += sph_lens.M_einasto(*lens_mass_para_i)
        return proj_mass
    
    def projected_mass_d1(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in radian
        xi = theta_rad*self.Dd          # radius in kpc
        dxi_proj_mass = 0.0             # mass gradient in Msun/kpc
        sph_lens = sph(xi,dim="2D")
        for i,lens_mass in enumerate(self.mass_model_list):
            lens_mass_para_i = self.mass_model_para[i]
            if lens_mass == "power_law":
                dxi_proj_mass += sph_lens.dM_power_law(*lens_mass_para_i)
            if lens_mass == "two_power_law":
                dxi_proj_mass += sph_lens.dM_two_power_law(*lens_mass_para_i)
            if lens_mass == "sersic":
                dxi_proj_mass += sph_lens.dM_sersic(*lens_mass_para_i)
            #if lens_mass == "einasto":
            #    d_proj_mass += sph_lens.M_einasto(*lens_mass_para_i)
        drad_proj_mass = dxi_proj_mass*self.Dd     # in unit Msun/rad 
        #d_proj_mass    = drad_proj_mass*arcsec2rad    # in unit Msun/arcsec
        return drad_proj_mass
    
    def projected_mass_d2(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in radian
        xi = theta_rad*self.Dd          # radius in kpc
        ddxi_proj_mass = 0.0            # mass 2nd derivative in Msun/kpc^2 
        sph_lens = sph(xi,dim="2D")
        for i,lens_mass in enumerate(self.mass_model_list):
            lens_mass_para_i = self.mass_model_para[i]
            if lens_mass == "power_law":
                ddxi_proj_mass += 2*np.pi*(sph_lens.power_law(*lens_mass_para_i)+xi*sph_lens.d_power_law(*lens_mass_para_i))
            if lens_mass == "two_power_law":
                ddxi_proj_mass += 2*np.pi*(sph_lens.two_power_law(*lens_mass_para_i)+xi*sph_lens.d_two_power_law(*lens_mass_para_i))
            if lens_mass == "sersic":
                ddxi_proj_mass += 2*np.pi*(sph_lens.sersic(*lens_mass_para_i)+xi*sph_lens.d_sersic(*lens_mass_para_i))
        ddrad_proj_mass = ddxi_proj_mass*self.Dd*self.Dd    # mass 2nd derivative in Msun/rad^2
        return ddrad_proj_mass
    
    def deflection_angle(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in unit of radian
        alpha_rad = 4.0*G*self.D_ratio/c**2*self.projected_mass(theta)/theta_rad
        alpha = alpha_rad/arcsec2rad    # deflection angle in arcsec
        return alpha
    
    def deflection_angle_d1(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in unit of radian
        drad_alpha_rad = 4.0*G*self.D_ratio/c**2/theta_rad*(self.projected_mass_d1(theta)-self.projected_mass(theta)/theta_rad)
        d_alpha = drad_alpha_rad    # d_alpha/d_theta, dimensionless
        return d_alpha 
    
    def deflection_angle_d2(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in unit of radian
        ddrad_alpha_rad = 4.0*G*self.D_ratio/c**2/theta_rad*(self.projected_mass_d2(theta)-2*self.projected_mass_d1(theta)/theta_rad+2*self.projected_mass(theta)/theta_rad**2)
        dd_alpha = ddrad_alpha_rad*arcsec2rad    # d^2_alpha/d_theta^2 in arcsec^-1
        return dd_alpha
    
    def convergence(self,theta):
        theta_rad = theta*arcsec2rad    # angular radius in unit of radian
        xi = theta_rad*self.Dd          # radius in kpc
        kappa = 0.0
        sph_lens = sph(xi,dim="2D")
        for i,lens_mass in enumerate(self.mass_model_list):
            lens_mass_para_i = self.mass_model_para[i]
            if lens_mass == "power_law":
                kappa += sph_lens.power_law(*lens_mass_para_i)/self.sigma_crit
            if lens_mass == "two_power_law":
                kappa += sph_lens.two_power_law(*lens_mass_para_i)/self.sigma_crit
            if lens_mass == "sersic":
                kappa += sph_lens.sersic(*lens_mass_para_i)/self.sigma_crit
            if lens_mass == "einasto":
                kappa += sph_lens.einasto(*lens_mass_para_i)/self.sigma_crit
        return kappa
    
    def lensing_equation(self,theta,beta = 0.0):
        return theta - beta - self.deflection_angle(theta)
            